package org.nhindirect.stagent.cert.impl;

import java.io.File;
import java.io.FileInputStream;
import java.io.FileOutputStream;
import java.security.Key;
import java.security.KeyStore;
import java.security.PrivateKey;
import java.security.cert.Certificate;
import java.security.cert.X509Certificate;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Enumeration;
import java.util.HashSet;
import java.util.Set;

import org.nhindirect.stagent.cert.CertificateStore;
import org.nhindirect.stagent.cert.X509CertificateEx;

public class KeyStoreCertificateService extends CertificateStore 
{
	private Set<X509Certificate> certs = new HashSet<X509Certificate>();
	
	/*
	 * TODO: change the way the passwords and the keystore are held
	 */
	private File keyStoreFile;
	private String keyStorePassword;
	private String privateKeyPassword;
	private KeyStore ks;
	
	public KeyStoreCertificateService()
	{
	}
	
	public KeyStoreCertificateService(File keyStoreFile)
	{
		this(keyStoreFile, null, null);
	}

	public KeyStoreCertificateService(File keyStoreFile, String keyStorePassword)
	{
		this(keyStoreFile, keyStorePassword, null);
	}
	
	public KeyStoreCertificateService(File keyStoreFile, String keyStorePassword, String privateKeyPassword)
	{
		this.keyStoreFile = keyStoreFile;
		this.keyStorePassword = keyStorePassword;
		this.privateKeyPassword = privateKeyPassword;
		
		if (keyStoreFile == null)
    		throw new IllegalArgumentException();

		bootstrapFromFile();
	}
	
	public void setKeyStoreFile(File fl)
	{
		if (ks != null)
			throw new IllegalStateException();			
		
		keyStoreFile = fl;
	}
	
	public void setKeyStoreFile(String fileName)
	{
		if (ks != null)
			throw new IllegalStateException();		
		
		setKeyStoreFile(new File(fileName));
	}
	
	public void setKeyStorePassword(String password)
	{
		if (ks != null)
			throw new IllegalStateException();
		
		this.keyStorePassword = password;
	}
	
	public void setPrivateKeyPassword(String privateKeyPassword)
	{
		if (ks != null)
			throw new IllegalStateException();
		
		this.privateKeyPassword = privateKeyPassword;
	}
	
	public void loadKeyStore()
	{
		if (ks != null)
			throw new IllegalStateException();
		
		bootstrapFromFile();
	}
	
	private void bootstrapFromFile()
	{
		try
		{
			ks = KeyStore.getInstance(KeyStore.getDefaultType());
			
			System.out.println("Loading certs from keystore file " + keyStoreFile.getAbsolutePath());
			 
			if (!keyStoreFile.exists())
			{
				// create a new keystore file
				ks.load(null, keyStorePassword == null ? null : keyStorePassword.toCharArray());
				
				FileOutputStream outStream = new FileOutputStream(keyStoreFile);
				ks.store(outStream, keyStorePassword == null ? null : keyStorePassword.toCharArray());
			}
			else
			{
				// load from keystore file
				FileInputStream inStream = new FileInputStream(keyStoreFile);
				
				ks.load(inStream, keyStorePassword == null ? null : keyStorePassword.toCharArray());
				
				Enumeration<String> aliases = ks.aliases();
				
				while (aliases.hasMoreElements())
				{
					String alias = aliases.nextElement();
					
					Certificate cert = ks.getCertificate(alias);
					if (cert != null && cert instanceof X509Certificate)
					{
						X509Certificate addCert;
						
						// check if there is private key
						Key key = ks.getKey(alias, privateKeyPassword == null ? null : privateKeyPassword.toCharArray());
						if (key != null && key instanceof PrivateKey)
							addCert = X509CertificateEx.fromX509Certificate((X509Certificate)cert, (PrivateKey)key);
						else
							addCert = (X509Certificate)cert;
						
						certs.add(addCert);
					}
				}
			}
		}
		catch (Exception e)
		{
			e.printStackTrace();
		}
	}
	
    public boolean contains(X509Certificate cert)
    {
    	return certs.contains(cert);
    }
    
    public void add(X509Certificate cert)
    {
		String newAlias = cert.getIssuerX500Principal().getName() + ":" + cert.getIssuerX500Principal().getName();   									

		add(cert, newAlias);
    }
    
    public void remove(X509Certificate cert)
    {
    	if (certs.remove(cert))
    	{
    		// remove from the key store
    		try
    		{
    			String alias = ks.getCertificateAlias(cert);
    			if (alias != null)
    			{
    				ks.deleteEntry(alias);
    				// persist
    				FileOutputStream outStream = new FileOutputStream(keyStoreFile);
    				ks.store(outStream, keyStorePassword == null ? null : keyStorePassword.toCharArray());	
    			}
    		}
    		catch (Exception e)
    		{
    			
    		}
    	}
    }
    
    public Collection<X509Certificate> getCertificates()
    {
    	// internal operations may iterate the returned collection for modifiable operations, so create a copy of the list
    	// instead of an unmodifiable collection backed by the set
    	return new ArrayList<X509Certificate>(certs);
    }
    
    public void add(X509Certificate cert, String alias)
    {
    	if (certs.contains(cert))
    		throw new IllegalArgumentException("Cert already contained in store.  Use update() to update a certificate");
    	
    	certs.add(cert);
    	
		try
		{
			
			if (cert instanceof X509CertificateEx)
				ks.setKeyEntry(alias, ((X509CertificateEx)cert).getPrivateKey(),
						privateKeyPassword == null ? null : privateKeyPassword.toCharArray(), new Certificate[] {cert});
			else
				ks.setCertificateEntry(alias, cert);
			
			// persist
			FileOutputStream outStream = new FileOutputStream(keyStoreFile);
			ks.store(outStream, keyStorePassword == null ? null : keyStorePassword.toCharArray());	
		}
		catch (Exception e)
		{
			
		}    	
    }
    
    public void update(X509Certificate cert, String alias)
    {
        if (contains(cert))
        {
            remove(cert);
        }
        add(cert, alias);
    }
    
    public X509Certificate getByAlias(String alias)
    {
    	X509Certificate retVal = null;
    	Certificate cert = null;
    	
    	try
    	{
    		cert = ks.getCertificate(alias);
			if (cert != null && cert instanceof X509Certificate)
			{
				X509Certificate addCert;
				
				// check if there is private key
				Key key = ks.getKey(alias, privateKeyPassword == null ? null : privateKeyPassword.toCharArray());
				if (key != null && key instanceof PrivateKey)
					retVal = X509CertificateEx.fromX509Certificate((X509Certificate)cert, (PrivateKey)key);
				else
					retVal = (X509Certificate)cert;				
			}    		    		
    	}
    	catch (Exception e)
    	{
    	}
    	
    	return retVal;
    }

}
